import os

from sklearn.model_selection import train_test_split

from option import parse_args


def save_file(video_names, original_path, file_name):
    video_names.sort()
    with open(os.path.join(original_path, file_name), 'w') as f:
        for video_name in video_names:
            f.write(video_name + '\n')


def main():
    args = parse_args()
    assert os.path.exists(args.root_dir)
    original_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', 'c23')
    original_videos_path = os.path.join(original_path, 'videos')
    video_names = os.listdir(original_videos_path)
    # shuffle
    # train_video_names, test_video_names = train_test_split(video_names, test_size=0.28, random_state=111)
    # test_video_names, val_video_names = train_test_split(test_video_names, test_size=0.5, random_state=111)
    # no shuffle
    video_names.sort()
    train_video_names, test_video_names = train_test_split(video_names, test_size=0.28, shuffle=False)
    val_video_names, test_video_names = train_test_split(test_video_names, test_size=0.5, shuffle=False)
    save_file(train_video_names, args.root_dir, 'train.txt')
    save_file(val_video_names, args.root_dir, 'val.txt')
    save_file(test_video_names, args.root_dir, 'test.txt')


if __name__ == '__main__':
    # Total: 2000
    main()
